%% Aim
% Calculate pupil information for each trial based on pupil.csv

%% Setup
tic;
clc;
clear;
root_path = pwd; 
addpath(genpath(root_path));

nrole=2;
nblock=2;
ntrial=100;
nprice=10;
nlottery=10;

%% Load data
beh_data=readtable([root_path,'\data\behavior\behavior.csv']);
pupil_data=readtable([root_path,'\data\pupil\pupil.csv']);
sub_data=readtable([root_path,'\data\subject\subject.csv']);
nsub=height(sub_data);

%% basic columns
Subject = beh_data.Subject;
Block = beh_data.Block;
Role = beh_data.Role;
PricePosition = beh_data.PricePosition;
Price = beh_data.Price;
Lottery = beh_data.Lottery;
Choice = beh_data.Choice;
RT = beh_data.RT;

%% Calculate baseline pupil, decision pupil (-500 to 1500 ms around decision, yet cut -750 to 1000 ms for illustration)

BasePupil = zeros(length(Price),1);
DecPupil = zeros(length(Price),1);

for isub=1:nsub
    for irole=1:2
        % import data
        pdata = pupil_data(pupil_data.Subject == isub & pupil_data.Role == irole,:);
        diffs=diff(pdata.StimulusName);
        events = find(diffs ~= 0) + 1;
        % loop through trials
        for i=1:length(events)
            if pdata.StimulusName(events(i))<=100
               continue;
            end
            % Cut out the trial
            option_on = pdata.TimeSignal(events(i));
            option_off = pdata.TimeSignal(events(i+1));
            pre_time = option_on - 500;
            predec_time = option_off - 750;            
            postdec_time = option_off + 1500;
            pre_option = pdata(pdata.TimeSignal >= pre_time & pdata.TimeSignal < option_on,:);
            whole_option = pdata(pdata.TimeSignal >= pre_time & pdata.TimeSignal <= postdec_time,:);
            whole_option_ex = pdata(pdata.TimeSignal >= pre_time & pdata.TimeSignal <= postdec_time+1500,:);
            whole_option_4000 = pdata(pdata.TimeSignal >= pre_time & pdata.TimeSignal <= option_on + 4000,:);
            % Baseline pupil
            base_pupil=mean(pre_option.Pupilz);
            % Baseline correction
            whole_option.Pupilz=whole_option.Pupilz-base_pupil;
            whole_option_ex.Pupilz=whole_option_ex.Pupilz-base_pupil;
            whole_option_4000.Pupilz=whole_option_4000.Pupilz-base_pupil;
            % Decision pupil window
            dec_option = pdata(pdata.TimeSignal >= predec_time & pdata.TimeSignal <= postdec_time,:);
            dec_option.Pupilz = dec_option.Pupilz-base_pupil;
            dec_pupil=mean(dec_option(dec_option.TimeSignal>=option_off - 500 & dec_option.TimeSignal>=option_on,:).Pupilz);
            
            % assign the value of variables
            price = floor(pdata.StimulusName(events(i))/1000);
            lottery = mod(pdata.StimulusName(events(i)),1000);
            idx = (Subject == isub & ...
                  Role == irole & ...
                  Price == price & ...
                  Lottery == lottery);
            BasePupil(idx) = base_pupil;
            DecPupil(idx) = dec_pupil;
            
            % save sample for ploting
            Pupil.Role(irole).PreOption{(price),(lottery)./2}=pre_option;
            Pupil.Role(irole).WholeOption{(price),(lottery)./2}=whole_option;
            Pupil.Role(irole).DecOption{(price),(lottery)./2}=dec_option;
            Pupil.Role(irole).WholeOptionEx{(price),(lottery)./2}=whole_option_ex;
            Pupil.Role(irole).WholeOption4000{(price),(lottery)./2}=whole_option_4000;
        end
    end
    pupil_trial_curve.Subject(isub)=Pupil;
end

save([root_path, '\data\pupil\pupil_trial_curve.mat'], 'pupil_trial_curve');

%% create a table .csv for all subjects
pupil_trial_data = table(Subject, Block, Role, PricePosition, Price, Lottery, Choice, RT, ...
                         BasePupil, DecPupil);
writetable(pupil_trial_data, [root_path, '\data\pupil\pupil_trial.csv']);

%%
toc;
